package leetcode_801_900;

import com.sun.deploy.util.ArrayUtil;
import node.Node;

import java.util.*;

public class LeeCode_834 {
    public static void main(String[] args) {
        System.out.println(Arrays.toString(sumOfDistancesInTree(6, new int[][]{{0, 1}, {0, 2}, {2, 3}, {2, 4}, {2, 5}})));
        sumOfDistancesInTree(1, new int[][]{});
    }
    private static int[] sumOfDistancesInTree(int n, int[][] edges) {
        if (n == 1)
            return new int[]{0};
        Map<Integer, List<Integer>> map = new HashMap<>();
        for (int[] edge : edges) {
            int x = edge[0], y = edge[1];
            if (map.containsKey(x)){
                map.get(x).add(y);
            }else {
                List<Integer> list = new ArrayList<>();
                list.add(y);
                map.put(x, list);
            }
            if (map.containsKey(y)){
                map.get(y).add(x);
            }else {
                List<Integer> list = new ArrayList<>();
                list.add(x);
                map.put(y, list);
            }
        }
        int[][] dp = new int[n][3];
        dfs(map, dp, 0, -1);
        dp[0][2] = dp[0][1];
        for (Integer next : map.get(0)) {
            dfs2(map, dp, next, 0, n);
        }
        int[] ans = new int[n];
        for (int i = 0; i < n; i++) {
            ans[i] = dp[i][2];
        }
        return ans;
    }
    private static void dfs2(Map<Integer, List<Integer>> map, int[][] dp, int now, int fa, int n){
        List<Integer> nextS = map.get(now);
        dp[now][2] = dp[fa][2] - dp[now][0] * 2 + n;
        for (Integer next : nextS) {
            if (next == fa)
                continue;
            dfs2(map, dp, next, now, n);
        }
    }
    private static void dfs(Map<Integer, List<Integer>> map, int[][] dp, int now, int fa){
        List<Integer> nextS = map.get(now);
        if (nextS.size() == 1 && nextS.get(0) == fa){
            dp[now][0] = 1;
            return;
        }
        int c0 = 1, c1 = 0;
        for (Integer next : nextS) {
            if (next == fa)
                continue;
            dfs(map, dp, next, now);
            c0 += dp[next][0];
            c1 += dp[next][0] + dp[next][1];
        }
        dp[now][0] = c0;
        dp[now][1] = c1;
    }
}